from neo4j import GraphDatabase
import os
from tqdm import tqdm

class Neo4jBatchExecutor:
    def __init__(self, uri, auth, cypher_dir="data/cypher"):
        self.driver = GraphDatabase.driver(uri, auth=auth)
        self.cypher_dir = os.path.abspath(cypher_dir)
        # 按截图文件顺序调整执行流程
        self.execution_order = [
            "!clean.cypher",         # 清空数据库
            "!nodes.cypher",         # 创建节点
            "!indexes.cypher",       # 创建索引
            "!relationships.cypher"  # 建立关系
        ]

    def _parse_cypher_file(self, file_path):
        """精准解析截图中的Cypher文件格式"""
        with open(file_path, 'r', encoding='utf-8') as f:
            content = f.read()
            # 分割语句并过滤空行/注释（匹配截图中的分号结尾）
            return [
                q.strip() 
                for q in content.split(';') 
                if q.strip() and not q.strip().startswith("//")
            ]

    def _execute_transaction(self, queries):
        """带进度条的事务执行"""
        with self.driver.session(database="neo4j") as session:
            for query in tqdm(queries, desc="执行进度", unit="stmt"):
                try:
                    session.execute_write(lambda tx: tx.run(query))
                except Exception as e:
                    print(f"执行失败: {str(e)}\n问题语句: {query}")
                    raise

    def execute_files(self):
        """按截图文件顺序执行"""
        print(f"加载目录: {self.cypher_dir}")
        for filename in self.execution_order:
            file_path = os.path.join(self.cypher_dir, filename)
            if not os.path.exists(file_path):
                print(f"⚠ 文件缺失: {os.path.basename(file_path)}")
                continue
            
            print(f"\n▶ 开始执行 {filename}")
            queries = self._parse_cypher_file(file_path)
            try:
                self._execute_transaction(queries)
                print(f"✓ {filename} 执行成功")
            except Exception as e:
                print(f"✗ {filename} 执行失败: {str(e)}")
                break
                
    def __del__(self):
        self.driver.close()

if __name__ == "__main__":
    # 配置参数（与截图路径一致）
    executor = Neo4jBatchExecutor(
        uri="bolt://localhost:7687",
        auth=("neo4j", "123456789"),
        cypher_dir="data/cypher"  # 匹配截图中的路径
    )
    executor.execute_files()